package io.github.gszj2018.jp_analysis;

import io.github.gszj2018.jp_analysis.target.Field;
import io.github.gszj2018.jp_analysis.target.Heap;
import io.github.gszj2018.jp_analysis.target.Target;
import io.github.gszj2018.jp_analysis.target.Variable;
import soot.*;
import soot.jimple.*;
import soot.jimple.toolkits.callgraph.ReachableMethods;
import soot.util.queue.QueueReader;

import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class ProgramTransformer extends SceneTransformer {
    private final AndersonSolver as;
    private final Set<String> mDone = new HashSet<>();

    private static final String PARAM_PREFIX = "@parameter";
    private static final String THIS_LOCAL = "@this";
    private static final String RET_LOCAL = "@return";
    private static final String ARRAY_PREFIX = "@element";
    private static final String ARRAY_FIELD_ALL = ARRAY_PREFIX + "_a";
    private static final String ARRAY_FIELD_UNKNOWN = ARRAY_PREFIX + "_u";
    private static final String EXCEPTION_FIELD = "@exception";
    private static final int STATIC_OBJ_ID = -1;

    private static final Target STATIC_OBJ = new Heap(STATIC_OBJ_ID);
    private static final Target THROWN_EXCEPTION = new Field(EXCEPTION_FIELD, STATIC_OBJ);

    private static final String BENCH_ALLOC = "<benchmark.internal.BenchmarkN: void alloc(int)>";
    private static final String BENCH_TEST = "<benchmark.internal.BenchmarkN: void test(int,java.lang.Object)>";

    public ProgramTransformer(AndersonSolver as) {
        this.as = as;
    }

    private static String param(int id) {
        return PARAM_PREFIX + id;
    }

    private static String element(int id) {
        return ARRAY_PREFIX + id;
    }

    @Override
    protected void internalTransform(String phaseName, Map<String, String> options) {
        ReachableMethods rm = Scene.v().getReachableMethods();
        QueueReader<MethodOrMethodContext> q = rm.listener();
        while (q.hasNext()) {
            SootMethod cur = q.next().method();
//            if (cur.isJavaLibraryMethod()) continue;

            boolean isNonVirtual = cur.isConstructor()
                    || cur.isStatic()
                    || cur.isStaticInitializer()
                    || cur.isPrivate();
            if (!isNonVirtual) checkMethodOverride(cur);

            if (cur.hasActiveBody())
                transformMethod(cur.getActiveBody(), cur.getSignature());
        }

    }

    private void checkMethodOverride(SootMethod m) {
        if (mDone.add(m.getSignature()))
            checkSuperclass(m, m.getDeclaringClass());
    }

    private void checkSuperclass(SootMethod m, SootClass c) {
        if (c.hasSuperclass())
            processMethodOverride(m, c.getSuperclass());

        for (SootClass si : c.getInterfaces())
            processMethodOverride(m, si);
    }

    private void processMethodOverride(SootMethod m, SootClass c) {
        SootMethod mBase = c.getMethodUnsafe(m.getSubSignature());
        if (mBase != null) {
            String base = mBase.getSignature();
            String derived = m.getSignature();
            // ret
            Target lhs = new Variable(base, RET_LOCAL);
            Target rhs = new Variable(derived, RET_LOCAL);
            as.addConstraint(lhs, rhs);
            // this
            lhs = new Variable(derived, THIS_LOCAL);
            rhs = new Variable(base, THIS_LOCAL);
            as.addConstraint(lhs, rhs);
            // args
            int ac = m.getParameterCount();
            for (int i = 0; i < ac; ++i) {
                lhs = new Variable(derived, param(i));
                rhs = new Variable(base, param(i));
                as.addConstraint(lhs, rhs);
            }
            // Success, process next link
            checkMethodOverride(mBase);
        } else {
            // Failed, try ancestors
            checkSuperclass(m, c);
        }
    }

    private void transformMethod(Body mBody, String mSig) {
        int allocId;
        int nextAllocId = 0;
        for (Unit u : mBody.getUnits()) {
            allocId = nextAllocId;
            nextAllocId = 0;
            if (u instanceof InvokeStmt) {
                InvokeExpr ie = ((InvokeStmt) u).getInvokeExpr();
                String iSig = ie.getMethod().getSignature();
                if (BENCH_ALLOC.equals(iSig)) {
                    nextAllocId = ((IntConstant) ie.getArg(0)).value;
                } else if (BENCH_TEST.equals(iSig)) {
                    int id = ((IntConstant) ie.getArg(0)).value;
                    as.addQuery(id, transformValue(ie.getArg(1), mSig, allocId));
                } else {
                    transformValue(ie, mSig, allocId);
                }
            } else if (u instanceof DefinitionStmt) {
                DefinitionStmt du = (DefinitionStmt) u;
                Target lhs = transformValue(du.getLeftOp(), mSig, allocId);
                Target rhs = transformValue(du.getRightOp(), mSig, allocId);
                if (lhs instanceof Field) {
                    // in case of array write
                    // write to UNKNOWN rather than ALL
                    // see transformValue() for more information
                    Field lf = (Field) lhs;
                    if (lf.fn.equals(ARRAY_FIELD_ALL)) {
                        lhs = new Field(ARRAY_FIELD_UNKNOWN, lf.o);
                    }
                }
                as.addConstraint(lhs, rhs);
            } else if (u instanceof ReturnStmt) {
                Target rhs = transformValue(((ReturnStmt) u).getOp(), mSig, allocId);
                as.addConstraint(new Variable(mSig, RET_LOCAL), rhs);
            } else if (u instanceof ThrowStmt) {
                Target rhs = transformValue(((ThrowStmt) u).getOp(), mSig, allocId);
                as.addConstraint(THROWN_EXCEPTION, rhs);
            }
        }
    }

    private Target transformValue(Value v, String mSig, int allocId) {
        if (v instanceof AnyNewExpr) {
            return new Heap(allocId);
        } else if (v instanceof CastExpr) {
            return transformValue(((CastExpr) v).getOp(), mSig, allocId);
        } else if (v instanceof InstanceFieldRef) {
            SootField f = ((InstanceFieldRef) v).getField();
            Value o = ((InstanceFieldRef) v).getBase();
            return new Field(f.getSignature(), transformValue(o, mSig, allocId));
        } else if (v instanceof StaticFieldRef) {
            SootField f = ((StaticFieldRef) v).getField();
            return new Field(f.getSignature(), STATIC_OBJ);
        } else if (v instanceof InvokeExpr) {
            InvokeExpr iv = (InvokeExpr) v;
            List<Value> args = iv.getArgs();
            String iSig = iv.getMethod().getSignature();
            for (int i = 0; i < args.size(); ++i) {
                Variable dst = new Variable(iSig, param(i));
                Target src = transformValue(args.get(i), mSig, allocId);
                as.addConstraint(dst, src);
            }
            if (iv instanceof InstanceInvokeExpr) {
                Variable dst = new Variable(iSig, THIS_LOCAL);
                Target src = transformValue(((InstanceInvokeExpr) iv).getBase(), mSig, allocId);
                as.addConstraint(dst, src);
            }
            return new Variable(iSig, RET_LOCAL);
        } else if (v instanceof Local) {
            return new Variable(mSig, ((Local) v).getName());
        } else if (v instanceof ParameterRef) {
            return new Variable(mSig, param(((ParameterRef) v).getIndex()));
        } else if (v instanceof ThisRef) {
            return new Variable(mSig, THIS_LOCAL);
        } else if (v instanceof ArrayRef) {
            // Array sensitivity for constant indexes
            // ======================================
            // R <----  ANY  <---
            //           ^      |
            //           |      |
            // RW <--> <idx>    |
            //           ^      |
            //           |      |
            // W ---> UNKNOWN ---
            // ======================================
            ArrayRef va = (ArrayRef) v;
            Target o = transformValue(va.getBase(), mSig, allocId);
            Value idx = va.getIndex();
            Target anyElement = new Field(ARRAY_FIELD_ALL, o);
            Target unknownElement = new Field(ARRAY_FIELD_UNKNOWN, o);
            if (idx instanceof IntConstant) {
                Target indexed = new Field(element(((IntConstant) idx).value), o);
                as.addConstraint(indexed, unknownElement);
                as.addConstraint(anyElement, indexed);
                return indexed;
            } else {
                as.addConstraint(anyElement, unknownElement);
                // for array writes, this will be replaced with unknownElement
                // see transformMethod()
                return anyElement;
            }
        } else if (v instanceof CaughtExceptionRef) {
            return THROWN_EXCEPTION;
        } else {
            return Target.UNSUPPORTED;
        }
    }
}
